#!/usr/bin/env python3
"""
upload_to_multi_git.py

工具说明:
- 扫描指定目录，按大小将文件分配到多个本地 git 仓库（每个仓库不超过指定阈值），
- 支持将每个仓库设置远端并分批次提交(push)，防止一次性提交过多文件导致卡死，
- 将本地文件路径与线上 raw 地址的映射记录到 sqlite 数据库，方便检索。
- 支持指定上传到Git仓库的子目录，而不是直接上传到根目录

使用前准备:
- Python 3.8+
- 安装依赖: `pip install -r requirements.txt`
- 如果需要自动创建 GitHub 仓库, 提供 `--github-token`。
- 如果要推送到现有远端，请提供一个包含远端仓库 URL 列表的文本文件 `--remotes-file`，每行一个 URL。

注: 脚本假定你本机已配置好 Git 凭据 (ssh 或 https credential helper)，或者使用 GitHub token 创建并推送。
"""

import argparse
import os
import sqlite3
import tempfile
import shutil
import sys
import json
from pathlib import Path
from typing import List, Tuple, Dict, Optional

try:
    from git import Repo
except Exception:
    print("请先安装依赖: pip install GitPython requests tqdm")
    raise

import requests
from tqdm import tqdm

import time


DB_SCHEMA = """
CREATE TABLE IF NOT EXISTS repos (
    id INTEGER PRIMARY KEY,
    name TEXT,
    remote_url TEXT,
    raw_base TEXT,
    branch TEXT,
    target_dir TEXT  -- 添加目标子目录字段
);
CREATE TABLE IF NOT EXISTS files (
    id INTEGER PRIMARY KEY,
    local_path TEXT,
    repo_id INTEGER,
    remote_path TEXT,  -- 修改为完整的远程路径（包含子目录）
    raw_url TEXT,
    FOREIGN KEY(repo_id) REFERENCES repos(id)
);
"""


def scan_files(root: str) -> List[Tuple[str, int]]:
    """返回 (relative_path, size_bytes) 列表"""
    files = []
    root = os.path.abspath(root)
    for dirpath, _, filenames in os.walk(root):
        for fn in filenames:
            full = os.path.join(dirpath, fn)
            try:
                st = os.stat(full)
            except OSError:
                continue
            rel = os.path.relpath(full, root)
            files.append((rel, st.st_size))
    files.sort(key=lambda x: x[1], reverse=True)
    return files


def partition_files(files: List[Tuple[str, int]], max_repo_bytes: int, max_files: Optional[int] = None) -> List[List[Tuple[str,int]]]:
    """简单贪心打包：大文件先放入当前包，超过阈值则开启新包"""
    packs: List[List[Tuple[str,int]]] = []
    cur_pack: List[Tuple[str,int]] = []
    cur_size = 0
    for rel, size in files:
        if max_files and len(cur_pack) >= max_files:
            packs.append(cur_pack)
            cur_pack = []
            cur_size = 0
        if size > max_repo_bytes:
            # 单文件超过阈值，单独成包（用户要注意）
            packs.append([(rel, size)])
            continue
        if cur_size + size > max_repo_bytes and cur_pack:
            packs.append(cur_pack)
            cur_pack = []
            cur_size = 0
        cur_pack.append((rel, size))
        cur_size += size
    if cur_pack:
        packs.append(cur_pack)
    return packs


def init_repo_at(path: str) -> Repo:
    repo = Repo.init(path)
    return repo


def add_and_commit_in_batches(repo_path: str, root_src: str, files: List[str], batch_count: int = 500, target_dir: str = "") -> List[List[str]]:
    """把相对路径文件从 root_src 复制到 repo_path，然后分批 git add/commit。返回每个批次的文件列表。
    
    Args:
        repo_path: 本地git仓库路径
        root_src: 源文件根目录
        files: 要添加的文件相对路径列表
        batch_count: 每批次提交的文件数
        target_dir: Git仓库中的目标子目录
    """
    batches = []
    repo = Repo(repo_path)
    wdir = Path(repo_path)
    
    # 确保目标子目录格式正确（移除末尾斜杠，使用正斜杠）
    if target_dir:
        target_dir = target_dir.rstrip('/').rstrip('\\')
    
    for i in range(0, len(files), batch_count):
        batch = files[i:i+batch_count]
        for rel in batch:
            src = os.path.join(root_src, rel)
            # 如果指定了目标子目录，则将文件复制到该子目录下
            if target_dir:
                dst_rel_path = os.path.join(target_dir, rel)
            else:
                dst_rel_path = rel
            
            dst = wdir.joinpath(dst_rel_path)
            dst.parent.mkdir(parents=True, exist_ok=True)
            if os.path.exists(src):
                shutil.copy2(src, dst)
        
        # 添加文件到git（注意路径格式转换）
        files_to_add = []
        for rel in batch:
            if target_dir:
                add_path = os.path.join(target_dir, rel)
            else:
                add_path = rel
            # 转换为git兼容的路径格式（使用正斜杠）
            add_path = add_path.replace('\\', '/')
            files_to_add.append(add_path)
        
        repo.index.add(files_to_add)
        msg = f"Add batch {i//batch_count + 1} files {len(batch)}"
        repo.index.commit(msg)
        batches.append(batch)
    return batches


def push_repo(repo: Repo, remote_url: Optional[str], branch: str = 'main', batch_size: int = 10) -> None:
    """
    推送仓库到远程，支持分批推送以避免超时
    
    Args:
        repo: Git仓库对象
        remote_url: 远程仓库URL
        branch: 分支名
        batch_size: 每批推送的提交数量
    """
    if remote_url:
        # 确保远程仓库存在
        if 'origin' in [r.name for r in repo.remotes]:
            origin = repo.remotes.origin
            # 更新远程 URL（如果有变化）
            if origin.url != remote_url:
                origin.set_url(remote_url)
        else:
            origin = repo.create_remote('origin', remote_url)
        
        # 优化Git配置以避免超时
        try:
            # 增加HTTP缓冲区大小
            repo.git.config('http.postBuffer', '524288000')  # 500MB
            # 增加超时时间
            repo.git.config('http.lowSpeedLimit', '1000')
            repo.git.config('http.lowSpeedTime', '600')  # 10分钟
        except Exception as e:
            print(f"配置Git参数时出错: {e}")
        
        # 获取本地提交历史
        local_commits = list(repo.iter_commits())
        
        try:
            # 尝试获取远程分支信息
            try:
                repo.git.fetch('origin')
            except Exception:
                print("无法获取远程信息，可能是首次推送")
            
            # 检查远程是否有相同分支
            remote_branch_exists = False
            try:
                for ref in repo.git.ls_remote('--heads', 'origin').split('\n'):
                    if f'refs/heads/{branch}' in ref:
                        remote_branch_exists = True
                        break
            except Exception:
                print("无法检查远程分支，假设不存在")
                remote_branch_exists = False
            
            if remote_branch_exists:
                try:
                    # 获取远程分支的最新提交
                    remote_commit = repo.commit(f'origin/{branch}')
                    # 找出本地比远程多的提交
                    ahead_commits = []
                    for commit in repo.iter_commits(f'{branch}'):
                        if commit.hexsha == remote_commit.hexsha:
                            break
                        ahead_commits.append(commit)
                    
                    # 反转列表，使最早的提交在前
                    ahead_commits.reverse()
                    
                    if len(ahead_commits) > 0:
                        print(f"有 {len(ahead_commits)} 个提交需要推送")
                        
                        # 分批推送
                        for i in range(0, len(ahead_commits), batch_size):
                            batch = ahead_commits[i:i+batch_size]
                            if not batch:
                                continue
                            
                            # 获取要推送的最近一个提交的hash
                            commit_hash = batch[-1].hexsha
                            print(f"推送批次 {i//batch_size + 1}/{(len(ahead_commits) + batch_size - 1) // batch_size}")
                            
                            # 尝试推送这一批次
                            for retry in range(3):  # 最多重试3次
                                try:
                                    repo.git.push('origin', f'{commit_hash}:refs/heads/{branch}')
                                    print(f"  推送成功")
                                    break
                                except Exception as push_error:
                                    print(f"  推送失败 (尝试 {retry + 1}/3): {push_error}")
                                    if retry < 2:
                                        print("  等待5秒后重试...")
                                        time.sleep(5)
                                    else:
                                        raise
                    else:
                        print("没有新的提交需要推送")
                except Exception as sync_error:
                    print(f"同步远程分支失败: {sync_error}")
                    print("尝试强制推送...")
                    # 如果同步失败，尝试强制推送
                    for retry in range(3):
                        try:
                            repo.git.push('--force', '--set-upstream', 'origin', branch)
                            print("已强制推送成功")
                            return
                        except Exception as force_error:
                            print(f"强制推送失败 (尝试 {retry + 1}/3): {force_error}")
                            if retry < 2:
                                print("等待5秒后重试...")
                                time.sleep(5)
                            else:
                                raise
            else:
                # 远程分支不存在，可能是首次推送
                print(f"远程分支 {branch} 不存在，准备首次推送")
                
                # 如果本地提交数量很多，也采用分批推送策略
                if len(local_commits) > batch_size:
                    print(f"本地有 {len(local_commits)} 个提交，将分批推送")
                    
                    # 先推送第一个提交建立分支
                    first_commit = local_commits[-1].hexsha
                    print("推送初始提交建立分支")
                    for retry in range(3):
                        try:
                            repo.git.push('--set-upstream', 'origin', f'{first_commit}:refs/heads/{branch}')
                            print("初始提交推送成功")
                            break
                        except Exception as first_push_error:
                            print(f"初始提交推送失败 (尝试 {retry + 1}/3): {first_push_error}")
                            if retry < 2:
                                print("等待5秒后重试...")
                                time.sleep(5)
                            else:
                                raise
                    
                    # 然后分批推送剩余提交
                    remaining_commits = local_commits[:-1]
                    remaining_commits.reverse()
                    
                    for i in range(0, len(remaining_commits), batch_size):
                        batch = remaining_commits[i:i+batch_size]
                        if not batch:
                            continue
                        
                        commit_hash = batch[-1].hexsha
                        print(f"推送批次 {i//batch_size + 1}/{(len(remaining_commits) + batch_size - 1) // batch_size}")
                        
                        for retry in range(3):
                            try:
                                repo.git.push('origin', f'{commit_hash}:refs/heads/{branch}')
                                print(f"  推送成功")
                                break
                            except Exception as batch_error:
                                print(f"  推送失败 (尝试 {retry + 1}/3): {batch_error}")
                                if retry < 2:
                                    print("  等待5秒后重试...")
                                    time.sleep(5)
                                else:
                                    raise
                else:
                    # 本地提交不多，直接推送
                    for retry in range(3):
                        try:
                            print(f"直接推送所有 {len(local_commits)} 个提交")
                            repo.git.push('--set-upstream', 'origin', branch)
                            print("推送成功")
                            break
                        except Exception as push_error:
                            print(f"推送失败 (尝试 {retry + 1}/3): {push_error}")
                            if retry < 2:
                                print("等待5秒后重试...")
                                time.sleep(5)
                            else:
                                raise
        except Exception as e:
            # 通用错误处理
            print(f"推送过程中出错: {e}")
            # 最后尝试推送单个提交
            try:
                if local_commits:
                    print("尝试推送最后一个提交...")
                    repo.git.push('origin', f'{local_commits[0].hexsha}:refs/heads/{branch}')
                    print("推送最后一个提交成功")
                else:
                    # 如果没有提交，尝试推送空分支
                    print("尝试推送空分支...")
                    repo.git.push('--set-upstream', 'origin', branch)
            except Exception as last_error:
                print(f"所有推送尝试都失败: {last_error}")
                raise


def derive_github_raw_base(remote_url: str, branch: str = 'main') -> Optional[str]:
    # 支持 github.com 的 SSH 或 HTTPS 格式
    # https://github.com/owner/repo.git => https://raw.githubusercontent.com/owner/repo/branch/
    s = remote_url
    if s.endswith('.git'):
        s = s[:-4]
    if 'github.com' in s:
        if s.startswith('git@'):
            # git@github.com:owner/repo
            parts = s.split(':', 1)[1]
            owner_repo = parts.strip('/')
        else:
            # https://github.com/owner/repo
            owner_repo = s.split('github.com')[-1].strip('/\n')
        return f"https://raw.githubusercontent.com/{owner_repo}/{branch}/"
    return None


def derive_raw_base(remote_url: str, branch: str = 'main', target_dir: str = "") -> Optional[str]:
    """生成raw文件的基础URL
    
    Args:
        remote_url: 远程仓库URL
        branch: 分支名
        target_dir: 目标子目录
        
    Returns:
        raw文件的基础URL，包含目标子目录（如果有）
    """
    s = remote_url
    if s.endswith('.git'):
        s = s[:-4]
    
    # 构建基础raw URL
    base_url = None
    
    # github
    if 'github.com' in s:
        if s.startswith('git@'):
            parts = s.split(':', 1)[1]
            owner_repo = parts.strip('/')
        else:
            owner_repo = s.split('github.com')[-1].strip('/\n')
        base_url = f"https://raw.githubusercontent.com/{owner_repo}/{branch}/"
    # gitee
    elif 'gitee.com' in s:
        if s.startswith('git@'):
            parts = s.split(':', 1)[1]
            owner_repo = parts.strip('/')
        else:
            owner_repo = s.split('gitee.com')[-1].strip('/\n')
        base_url = f"https://gitee.com/{owner_repo}/raw/{branch}/"
    # gitcode
    elif 'gitcode' in s:
        if s.startswith('git@'):
            parts = s.split(':', 1)[1]
            owner_repo = parts.strip('/')
        else:
            # try split by domain
            owner_repo = s.split('/')[-2] + '/' + s.split('/')[-1]
        base_url = f"https://gitcode.net/{owner_repo}/raw/{branch}/"
    
    # 如果指定了目标子目录，将其添加到base_url中
    if base_url and target_dir:
        # 确保路径格式正确（使用正斜杠，移除开头/结尾斜杠）
        target_dir = target_dir.strip('/').strip('\\').replace('\\', '/')
        base_url = f"{base_url}{target_dir}/"
    
    return base_url


def create_gitee_repo(token: str, name: str, private: bool = True) -> str:
    """使用 Gitee API 创建仓库，返回 ssh_url 或 https_url"""
    url = 'https://gitee.com/api/v5/user/repos'
    params = {'access_token': token}
    data = {'name': name, 'private': 'true' if private else 'false'}
    r = requests.post(url, params=params, json=data)
    if r.status_code not in (201,):
        raise RuntimeError(f"创建 Gitee 仓库失败: {r.status_code} {r.text}")
    j = r.json()
    # 返回 ssh_url 或 html_url
    return j.get('ssh_url') or j.get('html_url')


def create_gitcode_repo(token: str, name: str, private: bool = True) -> str:
    """占位: GitCode 自动创建依赖具体平台 API。如果你希望我实现，请提供 gitcode API 文档或 token 类型说明。
    目前默认抛出异常。
    """
    raise NotImplementedError('GitCode 自动创建未实现，请提供 remotes-file 或手动在 GitCode 上创建仓库')


def create_github_repo(token: str, name: str, private: bool = True) -> str:
    headers = {"Authorization": f"token {token}", "Accept": "application/vnd.github.v3+json"}
    data = {"name": name, "private": private}
    r = requests.post('https://api.github.com/user/repos', headers=headers, json=data)
    if r.status_code not in (201,):
        raise RuntimeError(f"创建 GitHub 仓库失败: {r.status_code} {r.text}")
    return r.json()['ssh_url']


def ensure_branch(repo: Repo, branch: str = 'main'):
    try:
        if branch in repo.heads:
            repo.heads[branch].checkout()
        else:
            repo.git.checkout('-b', branch)
    except Exception:
        pass


def save_mappings(db_path: str, repo_info: Dict, file_list: List[str], raw_base: Optional[str], target_dir: str = ""):
    """保存仓库和文件映射信息到数据库
    
    Args:
        db_path: 数据库文件路径
        repo_info: 仓库信息字典
        file_list: 文件相对路径列表
        raw_base: raw文件基础URL
        target_dir: 目标子目录
    """
    conn = sqlite3.connect(db_path)
    cur = conn.cursor()
    cur.executescript(DB_SCHEMA)
    
    # 插入仓库信息，包括目标子目录
    cur.execute("INSERT INTO repos (name, remote_url, raw_base, branch, target_dir) VALUES (?,?,?,?,?)",
                (repo_info.get('name'), repo_info.get('remote'), raw_base, repo_info.get('branch', 'main'), target_dir))
    repo_id = cur.lastrowid
    
    tups = []
    # 确保目标子目录格式正确
    if target_dir:
        target_dir = target_dir.rstrip('/').rstrip('\\')
    
    for rel in file_list:
        # 计算远程路径（包含子目录）
        if target_dir:
            remote_path = os.path.join(target_dir, rel).replace('\\', '/')
        else:
            remote_path = rel.replace('\\', '/')
        
        # 计算raw URL
        raw_url = None
        if raw_base:
            # 构建完整的raw URL（确保路径分隔符正确）
            file_path = rel.replace('\\', '/')
            raw_url = f"{raw_base}{file_path}"
        
        tups.append((rel, repo_id, remote_path, raw_url))
    
    # 插入文件信息，包括远程路径
    cur.executemany("INSERT INTO files (local_path, repo_id, remote_path, raw_url) VALUES (?,?,?,?)", tups)
    conn.commit()
    conn.close()


def main():
    p = argparse.ArgumentParser()
    p.add_argument('src_dir', help='要上传的源目录')
    p.add_argument('--max-repo-mb', type=int, default=500, help='每个仓库最大容量(单位MB)，默认500MB')
    p.add_argument('--max-files-per-repo', type=int, default=0, help='每个仓库最大文件数(0表示不限制)')
    p.add_argument('--batch-size', type=int, default=500, help='每次 commit 的文件数，默认500')
    p.add_argument('--remotes-file', help='包含远端仓库 URL 的文件(每行一个)，如果提供，将按顺序使用这些远端')
    p.add_argument('--github-token', help='如果要自动创建 GitHub 仓库，请提供 token')
    p.add_argument('--gitee-token', help='如果要自动创建 Gitee 仓库，请提供 gitee access_token')
    p.add_argument('--gitcode-token', help='（可选）GitCode token（若支持自动创建）')
    p.add_argument('--create-platform', choices=['github','gitee','gitcode'], help='如果提供，将在指定平台自动创建仓库（需要相应 token）')
    p.add_argument('--repo-name-prefix', default='auto-split', help='自动创建时仓库名前缀')
    p.add_argument('--db', default='file_mappings.db', help='sqlite db 路径')
    p.add_argument('--work-dir', default=None, help='工作目录，默认使用临时目录')
    p.add_argument('--target-dir', default='', help='Git仓库中的目标子目录，默认为空（根目录）')
    args = p.parse_args()

    src = args.src_dir
    if not os.path.isdir(src):
        print('源目录不存在')
        sys.exit(1)

    files = scan_files(src)
    print(f'发现 {len(files)} 个文件，总体积 {sum(s for _, s in files)/1024/1024:.2f} MB')

    max_repo_bytes = args.max_repo_mb * 1024 * 1024
    max_files = args.max_files_per_repo if args.max_files_per_repo > 0 else None
    packs = partition_files(files, max_repo_bytes, max_files)
    print(f'将分为 {len(packs)} 个仓库（组）')

    remotes = []
    if args.remotes_file:
        with open(args.remotes_file, 'r', encoding='utf-8') as f:
            remotes = [line.strip() for line in f if line.strip()]

    work_root = args.work_dir or tempfile.mkdtemp(prefix='multi_git_')
    print('工作目录:', work_root)
    
    # 显示目标子目录信息
    if args.target_dir:
        print(f'上传目标子目录: {args.target_dir}')
    else:
        print('上传到仓库根目录')

    repo_count = 0
    for idx, pack in enumerate(packs):
        repo_count += 1
        name = f"{args.repo_name_prefix}-{idx+1}"
        repo_path = os.path.join(work_root, name)
        os.makedirs(repo_path, exist_ok=True)
        repo = init_repo_at(repo_path)
        ensure_branch(repo, 'main')

        # 准备文件列表
        rels = [rel for rel, _ in pack]
        print(f'组 {idx+1}: {len(rels)} 文件, 大小 {sum(s for _, s in pack)/1024/1024:.2f} MB -> 本地仓库 {repo_path}')

        # 逐批复制并 commit，传入目标子目录
        batches = add_and_commit_in_batches(repo_path, src, rels, batch_count=args.batch_size, target_dir=args.target_dir)

        # 处理远端：优先使用 remotes 列表；否则根据 create_platform 或可用 token 尝试创建仓库
        remote_url = None
        raw_base = None
        if remotes and idx < len(remotes):
            remote_url = remotes[idx]
        else:
            platform = args.create_platform
            try:
                if platform == 'github' and args.github_token:
                    remote_url = create_github_repo(args.github_token, name)
                elif platform == 'gitee' and args.gitee_token:
                    remote_url = create_gitee_repo(args.gitee_token, name)
                elif platform == 'gitcode' and args.gitcode_token:
                    remote_url = create_gitcode_repo(args.gitcode_token, name)
                else:
                    # 没有指定 create_platform，尝试按优先级使用提供的 token
                    if args.github_token:
                        remote_url = create_github_repo(args.github_token, name)
                    elif args.gitee_token:
                        remote_url = create_gitee_repo(args.gitee_token, name)
                    elif args.gitcode_token:
                        remote_url = create_gitcode_repo(args.gitcode_token, name)
            except NotImplementedError as nie:
                print(f'创建远端仓库未实现或不可用: {nie}')
            except Exception as e:
                print(f'创建远端仓库失败: {e}')

        if remote_url:
            try:
                push_repo(repo, remote_url, 'main')
            except Exception as e:
                print('推送异常，继续下一组:', e)

            # 推送后尝试推断 raw 地址（支持 github/gitee/gitcode），传入目标子目录
            raw_base = derive_raw_base(remote_url, 'main', args.target_dir)

        repo_info = {'name': name, 'remote': remote_url, 'branch': 'main'}
        # 保存映射（使用相对路径），传入目标子目录
        save_mappings(args.db, repo_info, rels, raw_base, args.target_dir)

    print('全部组处理完成。')
    print('数据库位置:', args.db)


if __name__ == '__main__':
    main()